#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
import json

def get_overlap_area(rect1, rect2):
    x1 = min(rect1[0], rect1[2])
    y1 = min(rect1[1], rect1[3])
    x2 = max(rect1[0], rect1[2])
    y2 = max(rect1[1], rect1[3])
    x3 = min(rect2[0], rect2[2])
    y3 = min(rect2[1], rect2[3])
    x4 = max(rect2[0], rect2[2])
    y4 = max(rect2[1], rect2[3])

    # 与标准矩形的一边相交
    if (x3 <= x1 and y3 >= y1 and y3 <= y2 and x4 >= x1 and x4 <= x2 and y4 >= y1 and y4 <= y2): # 左
        return (x4 - x1 + 1) * (y4 - y3 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 <= y1 and x4 >= x1 and x4 <= x2 and y4 >= y1 and y4 <= y2): # 上
        return (x4 - x3 + 1) * (y4 - y1 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 >= y1 and y3 <= y2 and x4 >= x2 and y4 >= y1 and y4 <= y2): # 右
        return (x2 - x3 + 1) * (y4 - y3 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 >= y1 and y3 <= y2 and x4 >= x1 and x4 <= x2 and y4 >= y2): # 下
        return (x4 - x3 + 1) * (y2 - y3 + 1)
    
    # 与标准矩形的相邻两边相交
    elif (x3 <= x1 and y3 <= y1 and x4 >= x1 and x4 <= x2 and y4 >= y1 and y4 <= y2): # // 左上
        return (x4 - x1 + 1) * (y4 - y1 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 <= y1 and x4 >= x2 and y4 >= y1 and y4 <= y2): # // 右上
        return (x2 - x3 + 1) * (y4 - y1 + 1)
    elif (x3 <= x1 and y3 >= y1 and y3 <= y2 and x4 >= x1 and x4 <= x2 and y4 >= y2): # // 左下
        return (x4 - x1 + 1) * (y2 - y3 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 >= y1 and y3 <= y2 and x4 >= x2 and y4 >= y2): # // 右下
        return (x2 - x3 + 1) * (y2 - y3 + 1)

    # 与标准矩形的平行两边相交
    elif (x3 <= x1 and y3 <= y1 and x4 >= x1 and x4 <= x2 and y4 >= y2): # // 上下，靠左
        return (x4 - x1 + 1) * (y2 - y1 + 1)
    elif (x3 <= x1 and y3 <= y1 and x4 >= x2 and y4 >= y1 and y4 <= y2): # // 左右，靠上
        return (x2 - x1 + 1) * (y4 - y1 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 <= y1 and x4 >= x2 and y4 >= y2): # // 上下，靠右
        return (x2 - x3 + 1) * (y2 - y1 + 1)
    elif (x3 <= x1 and y3 >= y1 and y3 <= y2 and x4 >= x2 and y4 >= y2): # // 左右，靠下
        return (x2 - x1 + 1) * (y2 - y3 + 1)
    
    # 贯穿标准上下或左右
    elif (x3 >= x1 and x3 <= x2 and y3 <= y1 and x4 >= x1 and x4 <= x2 and y4 >= y2): # // 上下
        return (x4 - x3 + 1) * (y2 - y1 + 1)
    elif (x3 <= x1 and y3 >= y1 and y3 <= y2 and x4 >= x2 and y4 >= y1 and y4 <= y2): # // 左右
        return (x2 - x1 + 1) * (y4 - y3 + 1)

    # 与标准矩形不相交但有重合
    elif (x3 <= x1 and y3 <= y1 and x4 >= x2 and y4 >= y2): # // 在标准答案外
        return (x2 - x1 + 1) * (y2 - y1 + 1)
    elif (x3 >= x1 and x3 <= x2 and y3 >= y1 and y3 <= y2 and x4 >= x1 and x4 <= x2 and y4 >= y1 and y4 <= y2): # // 在标准答案内
        return (x4 - x3 + 1) * (y4 - y3 + 1)
    else: #{ // 不相交也不重合
        return 0

    assert(False)
    return -1


def calculate_iou(gt, dt):
    gt_xmin, gt_ymin, gt_xmax, gt_ymax = gt[0], gt[1], gt[2], gt[3]
    dt_xmin, dt_ymin, dt_xmax, dt_ymax = dt[0], dt[1], dt[2], dt[3]

    # 获取矩形框交集对应的顶点坐标(intersection)
    xx1 = np.max([gt_xmin, dt_xmin])
    yy1 = np.max([gt_ymin, dt_ymin])
    xx2 = np.min([gt_xmax, dt_xmax])
    yy2 = np.min([gt_ymax, dt_ymax])

    # 计算两个矩形框面积
    area1 = int((gt_xmax - gt_xmin) * (gt_ymax - gt_ymin))
    area2 = int((dt_xmax - dt_xmin) * (dt_ymax - dt_ymin))
 
    # 计算交集面积 
    inter_area = int((np.max([0, xx2 - xx1])) * (np.max([0, yy2 - yy1 ])))
    #inter_area2 = get_overlap_area(gt, dt)
    #assert(int(inter_area)==int(inter_area2))
    
    #print(inter_area, inter_area2)
    #input("watch")
    # 计算交并比
    iou = inter_area / (area1 + area2 - inter_area)
    #print("iou1 method1", xx1, yy1, inter_area, area1 + area2 - inter_area, iou)
    return iou


def calculate_iou2(bbox1, bbox2):
    #xmin,ymin,xmax,ymax
    left = max(bbox1[0], bbox2[0])
    top = max(bbox1[1], bbox2[1])
    right = min(bbox1[2], bbox2[2])
    bottom = min(bbox1[3], bbox2[3])
    if left < right and top < bottom:
        intersection = (right - left) *  (bottom - top)
        b1Area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
        b2Area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
        iou = intersection * 1.0/ (b1Area + b2Area - intersection)
        #print("iou2 method1", left, top, intersection, b1Area + b2Area - intersection, iou)
        return iou
    else:
        return 0.0 


def get_dalei_categories(categories):
    dalei_categories = {}
    for cat_name in categories:
        # 跳过不参与计数的小类别
        if len(categories[cat_name]["supercategory"]) == 0:
            continue
        if categories[cat_name]["supercategory"] not in dalei_categories:
            dalei_categories[categories[cat_name]["supercategory"]] = {
                "gtn": 0, 
                "dtn": 0,
                "matched": 0,
                "subcategory": [],
                "thresh": [],
                "iou": [],
                "max_match": [],
                "weight": []
            }
        dalei_categories[categories[cat_name]["supercategory"]]["gtn"] += categories[cat_name]["gtn"]
        dalei_categories[categories[cat_name]["supercategory"]]["subcategory"].append(cat_name)
        dalei_categories[categories[cat_name]["supercategory"]]["thresh"].append(categories[cat_name]["thresh"])
        dalei_categories[categories[cat_name]["supercategory"]]["iou"].append(categories[cat_name]["iou"])
        dalei_categories[categories[cat_name]["supercategory"]]["max_match"].append(categories[cat_name]["max_match"])
        dalei_categories[categories[cat_name]["supercategory"]]["weight"].append(categories[cat_name]["weight"])

    for cat_name in dalei_categories:
        dalei_categories[cat_name]["thresh"] = np.mean(dalei_categories[cat_name]["thresh"])
        dalei_categories[cat_name]["iou"] = np.mean(dalei_categories[cat_name]["iou"])
        dalei_categories[cat_name]["max_match"] = np.mean(dalei_categories[cat_name]["max_match"])
        dalei_categories[cat_name]["weight"] = np.mean(dalei_categories[cat_name]["weight"])
        dalei_categories[cat_name]["del"] = 0

    return dalei_categories


def trans2dalei(categories, xiaolei):
    dalei = []
    for xiao in xiaolei:
        for cat_name in categories:
            if xiao["category"] not in categories[cat_name]["subcategory"]:
                continue
            xiao["category"] = cat_name
            dalei.append(xiao)
    
    return dalei


color_plane = {
    "ann_matched": (0,255,0),
    "ann_unmatched": (0, 0, 255),
    "det_matched": (255, 0, 0),
    "det_unmatched": (0, 0, 0)
}


def draw_objects(image, annos, dets, matched_pair, categories):
    
    matched_dets = []
    for ann_index, ann in enumerate(annos):
        if ann_index in matched_pair:  # ann_matched and det_matched
            det = dets[matched_pair[ann_index]]
            matched_dets.append(matched_pair[ann_index])
            assert(det["category"] == ann["category"])
            cv2.rectangle(image, (int(ann["bbox"][0]), int(ann["bbox"][1])), 
                          (int(ann["bbox"][2]), int(ann["bbox"][3])),
                          color_plane["ann_matched"], 3)

            cv2.rectangle(image, (int(det["bbox"][0]), int(det["bbox"][1])), 
                          (int(det["bbox"][2]), int(det["bbox"][3])), 
                           color_plane["det_matched"], 3)
            cv2.putText(image, det["category"]+"|"+"{:.2f}".format(det["score"]),
                        (int(det["bbox"][0]), int(det["bbox"][1])), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["det_matched"], 2)
        else:  # ann_unmatched
            cv2.rectangle(image, (int(ann["bbox"][0]), int(ann["bbox"][1])), 
                          (int(ann["bbox"][2]), int(ann["bbox"][3])), 
                           color_plane["ann_unmatched"], 3)

            cv2.putText(image, ann["category"], 
                        (int(ann["bbox"][0]), int(ann["bbox"][1])), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["ann_unmatched"], 2)

    # det_unmatched
    for det_index, det in enumerate(dets):
        if det_index in matched_dets:
            continue
        
        cv2.rectangle(image, (int(det["bbox"][0]), int(det["bbox"][1])), 
                        (int(det["bbox"][2]), int(det["bbox"][3])), 
                        color_plane["det_unmatched"], 3)
        cv2.putText(image, det["category"]+"|"+"{:.2f}".format(det["score"]), 
                    (int(det["bbox"][0]), int(det["bbox"][1])), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["det_unmatched"], 2)
    
    color_text_blue = "blue for det_matched"
    color_text_green = "gree for ann_matched"
    color_text_red = "red for ann_unmatched"
    color_text_black = "black for det_unmatched"
    # TODO: 自动获取文本高度
    cv2.putText(image, color_text_green, (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["ann_matched"], 2)
    cv2.putText(image, color_text_red, (20, 110), cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["ann_unmatched"], 2)
    cv2.putText(image, color_text_blue, (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["det_matched"], 2)
    cv2.putText(image, color_text_black, (20, 140), cv2.FONT_HERSHEY_SIMPLEX, 1.2, color_plane["det_unmatched"], 2)
    return image

image_name_map = {
    "施工吊车": "ShiGongDiaoChe",
    "500吨吊车": "500DunDiaoChe",
    "500吨吊车施工": "500DunDiaoCheShiGong",
    "SANY吊车": "SANYDiaoChe",
    "大马力吊车作业": "DaMaLiDiaoCheZuoYe",
    "大型吊车": "DaXingDiaoChe",
    "电杆吊装": "DianGanDiaoZhuang",
    "电力吊车": "DianLiDiaoChe",
    "电力吊车电缆吊运": "DianLiDiaoCheDianLanDiaoYun",
    "电力吊车起吊": "DianLiDiaoCheQiDiao",
    "电力吊车设备吊装": "DiaoLiDiaoCheSheBeiDiaoZhuang",
    "电力吊车施工": "DiaoLiDiaoCheShiGong",
    "电力施工": "DianLiShiGong",
    "电力作业": "DianLiZuoYe",
    "吊车": "DiaoChe",
    "吊车施工": "DiaoCheShiGong",
    "国网吊车": "GuoWangDiaoChe",
    "国家电网吊车作业": "GuoJiaDianWangDiaoCheZuoYe",
    "进口吊车": "JinKouDiaoChe",
    "汽车吊": "QiCheDiao",
    "三一吊车": "SanYiDiaoChe",
    "三一吊车作业": "SanYiDiaoCheZuoYe",
    "三一液压吊车作业": "SanYiYeYaDiaoCheZuoYe",
    "山东重工吊车": "ShanDongZhongGongDiaoChe",
    "山东重工液压吊车": "ShanDongZhongGongYeYaDiaoChe",
    "三一液压吊车": "SanYiYeYaDiaoChe",
    "山东重工液压吊车施工": "ShanDongZhongGongYeYaDiaoCheShiGong",
    "徐工起重吊车": "XuGongQiZhongDiaoChe",
    "液压吊车": "YeYaDiaoChe",
    "电力塔": "DianLiTa",
    "电力设备维护": "DianLiSheBeiWeiHu",
    "电力高架": "DianLiGaoJia",
    "电力高架线路": "DianLiGaoJiaXianLu",
    "电力高架施工": "DianLiGaoJiaShiGong",
    "电力高架监控": "DianLiGaoJiaJianKong",
    "三一吊车施工": "SanYiDiaoCheShiGong",
    "国网吊车施工": "GuoWangDiaoCheShiGong",
    "吊车电缆放电": "DiaoCheDianLanFangDian",
    "吊车电网施工": "DiaoCheDianWangShiGong",
    "工地钢筋工作业": "GongYeGangJinZuoYe"
}

def get_name_id_map(testCocoJson):
    file_id_map = {}
    id_file_map = {}
    id_categorie_map = {}
    with open(testCocoJson, "r", encoding='utf-8') as file:
    # with open(testCocoJson) as file:
        data = json.load(file)
    images = data["images"]
    for image in images:
        assert(image["id"] not in id_file_map)

        file_name = image["file_name"]
        if file_name.split("_")[0] not in image_name_map:
            print("no image name map: {}".format(file_name.split("_")[0]))
            continue
        new_file_name = file_name.replace(file_name.split("_")[0], image_name_map[file_name.split("_")[0]])
        id_file_map[image["id"]] = new_file_name[:-4]
        file_id_map[new_file_name[:-4]] = image["id"]
    categories = data["categories"]
    for cat in categories:
        id_categorie_map[cat["id"]] = cat["name"]
    return file_id_map,id_file_map,id_categorie_map